{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Action recognition using TAO ActionRecognitionNet\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## Sample prediction of ActionRecognitionNet\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/action_recognition/ARNet_inference.png?raw=true\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train 3D RGB only for action recognition on the subset of [HMDB51](https://serre-lab.clps.brown.edu/resource/hmdb-a-large-human-motion-database/) dataset.\n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .etlt file for deployment to DeepStream.\n",
    "\n",
    "At the end of this notebook, you will have a trained and optimized `action_recognition` model that you\n",
    "may deploy via [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of ActionRecognitionNet using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"actionrecognitionnet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"actionrecognitionnet\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/action_recognition_net\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "\n",
    "# Set your encryption key, and use the same key for all commands\n",
    "%env KEY = nvidia_tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the [HMDB51](https://serre-lab.clps.brown.edu/resource/hmdb-a-large-human-motion-database/) dataset for the tutorial. Download the HMDB51 dataset and unrar them firstly (We choose fall_floor/ride_bike for this tutorial): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install unrar\n",
    "# NOTE: The following commands require `sudo`. You can run the command outside the notebook.\n",
    "!apt update\n",
    "!apt-get install unrar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# download the dataset and unrar the files\n",
    "!wget --no-check-certificate -P $HOST_DATA_DIR http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/hmdb51_org.rar\n",
    "!mkdir -p $HOST_DATA_DIR/videos && unrar x -o+ $HOST_DATA_DIR/hmdb51_org.rar $HOST_DATA_DIR/videos\n",
    "!mkdir -p $HOST_DATA_DIR/raw_data\n",
    "!unrar x -o+ $HOST_DATA_DIR/videos/fall_floor.rar $HOST_DATA_DIR/raw_data\n",
    "!unrar x -o+ $HOST_DATA_DIR/videos/ride_bike.rar $HOST_DATA_DIR/raw_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clone the dataset process script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!if [ -d tao_toolkit_recipes ]; then rm -rf tao_toolkit_recipes; fi\n",
    "!git clone https://github.com/NVIDIA-AI-IOT/tao_toolkit_recipes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the dependency for data generator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install xmltodict opencv-python"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the process script. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB_RGB.sh $HOST_DATA_DIR/raw_data $HOST_DATA_DIR/processed_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also provide scripts to preprocess optical flow dataset. The following cells for processing optical flow dataset is `Optional`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`OPTIONAL:` Download the app based on NVOF SDK to generate optical flow. It is packaged with this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!echo <passwd> | sudo -S apt install -y libfreeimage-dev"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`OPTIONAL` Run the process script for HMDB. \n",
    "\n",
    "`IMPORTANT NOTE`: to run the `preprocess_HMDB.sh` generating optical flow, a Turing or Ampere above GPU is needed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!cp ./AppOFCuda tao_toolkit_recipes/tao_action_recognition/data_generation/\n",
    "#!cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB.sh $HOST_DATA_DIR/raw_data $HOST_DATA_DIR/processed_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# download the split files and unrar\n",
    "!wget --no-check-certificate -P $HOST_DATA_DIR http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/test_train_splits.rar\n",
    "!mkdir -p $HOST_DATA_DIR/splits && unrar x -o+ $HOST_DATA_DIR/test_train_splits.rar $HOST_DATA_DIR/splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run split_HMDB to generate training split\n",
    "!if [ -d $HOST_DATA_DIR/train ]; then rm -rf $HOST_DATA_DIR/train $HOST_DATA_DIR/test; fi\n",
    "!cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && python3 ./split_dataset.py $HOST_DATA_DIR/processed_data $HOST_DATA_DIR/splits/testTrainMulti_7030_splits $HOST_DATA_DIR/train  $HOST_DATA_DIR/test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# verify\n",
    "!ls -l $HOST_DATA_DIR/train\n",
    "!ls -l $HOST_DATA_DIR/train/ride_bike\n",
    "!ls -l $HOST_DATA_DIR/test\n",
    "!ls -l $HOST_DATA_DIR/test/ride_bike"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Download pretrained model from NGC"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to https://ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "import os\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/actionrecognitionnet:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC \n",
    "!ngc registry model download-version \"nvidia/tao/actionrecognitionnet:trainable_v1.0\" --dest $HOST_RESULTS_DIR/pretrained\n",
    "# Download the optical flow model from NGC\n",
    "# !ngc registry model download-version \"nvidia/tao/actionrecognitionnet:trainable_v2.0\" --dest $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained/actionrecognitionnet_vtrainable_v1.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* model: configure the model setting\n",
    "    * model_type: type of model, rgb/of/joint\n",
    "    * backbone: resnet18/34/50/101/152 \n",
    "    * rgb_seq_length: length of RGB input sequence\n",
    "    * input_type: 2d/3d\n",
    "    * sample_strategy: consecutive\n",
    "    * dropout_ratio: probability to drop the hidden units\n",
    "* train: configure the training hyperparameters\n",
    "    * optim_config\n",
    "    * epochs\n",
    "    * checkpoint_interval\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * train_dataset_dir\n",
    "    * val_dataset_dir\n",
    "    * label_map: map the class label to id\n",
    "    * output_shape\n",
    "    * batch_size\n",
    "    * workers: number of workers to do data loading\n",
    "    * clips_per_video: number of clips to be sampled from single video\n",
    "    * augmentation_config\n",
    "\n",
    "Please refer to the TAO documentation about ActionRecognitionNet to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/experiment_rgb_3d_finetune.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Train 3D model:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We provide pretrained RGB-only model trained on HMDB5 dataset. With the pretrained model, we can even get better accuracy with less epochs.\n",
    "\n",
    "`KNOWN ISSUE`: \n",
    "- 1) The training log will be corrupted by pytorch warning in the notebook:\n",
    "\n",
    "     `[W pthreadpool-cpp.cc:90] Warning: Leaking Caffe2 thread-pool after fork. (function pthreadpool)` \n",
    "     \n",
    "     To see the full log in std out, please run the command in terminal. \n",
    "- 2) \"=\" in the checkpoint file name should removed before using the checkpoint in command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train RGB only model with PTM\")\n",
    "!tao model action_recognition train \\\n",
    "                  -e $SPECS_DIR/experiment_rgb_3d_finetune.yaml \\\n",
    "                  encryption_key=$KEY \\\n",
    "                  results_dir=$RESULTS_DIR/rgb_3d_ptm \\\n",
    "                  model.rgb_pretrained_model_path=$RESULTS_DIR/pretrained/actionrecognitionnet_vtrainable_v1.0/resnet18_3d_rgb_hmdb5_32.tlt  \\\n",
    "                  model.rgb_pretrained_num_classes=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training command for multi-gpu training. We can define the number of gpus and specify which GPU's are to be used by setting the `train.gpu_ids` parameter.\n",
    "## The following command will trigger multi-gpu training on gpu 0 and gpu 1.\n",
    "# !tao model action_recognition train \\\n",
    "#                   -e $SPECS_DIR/experiment_rgb_3d_finetune.yaml \\\n",
    "#                   encryption_key=$KEY \\\n",
    "#                   train.gpu_ids=[0,1] \\\n",
    "#                   results_dir=$RESULTS_DIR/rgb_3d_ptm \\\n",
    "#                   model.rgb_pretrained_model_path=$RESULTS_DIR/pretrained/actionrecognitionnet_vtrainable_v1.0/resnet18_3d_rgb_hmdb5_32.tlt  \\\n",
    "#                   model.rgb_pretrained_num_classes=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Encrypted checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/rgb_3d_ptm/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/rgb_3d_ptm/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"rgb_3d_ptm/train/ar_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/rgb_3d_ptm/train/rgb_only_model.tlt\n",
    "!ls -ltrh $HOST_RESULTS_DIR/rgb_3d_ptm/train/rgb_only_model.tlt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `OPTIONAL` 4.2 Train optical flow only model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Important Note` The following cells are using optical flow dataset. \n",
    "\n",
    "We will train a 3D OF-only model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# print(\"Train 3D OF-only model\")\n",
    "# !tao model action_recognition train \\\n",
    "#                   -e $SPECS_DIR/experiment_of_3d_finetune.yaml \\\n",
    "#                   encryption_key=$KEY \\\n",
    "#                   results_dir=$RESULTS_DIR/of_3d_ptm \\\n",
    "#                   dataset.train_dataset_dir=$DATA_DIR/train \\\n",
    "#                   dataset.val_dataset_dir=$DATA_DIR/test \\\n",
    "#                   model.of_pretrained_model_path=$RESULTS_DIR/pretrained/actionrecognitionnet_vtrainable_v2.0/resnet18_3d_of_hmdb5_32_a100.tlt  \\\n",
    "#                   model.of_pretrained_num_classes=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"To resume training from a checkpoint, set the resume_training_checkpoint_path option to be the .tlt you want to resume from\")\n",
    "# print(\"remember to remove the `=` in the checkpoint's file name\")\n",
    "# !tao model action_recognition train \\\n",
    "#                   -e $SPECS_DIR/experiment_of_3d_finetune.yaml \\\n",
    "#                   encryption_key=$KEY \\\n",
    "#                   results_dir=$RESULTS_DIR/of_3d_ptm \\\n",
    "#                   train.resume_training_checkpoint_path="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print('Encrypted checkpoints:')\n",
    "# print('---------------------')\n",
    "# !ls -ltrh $HOST_RESULTS_DIR/of_3d_ptm/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/of_3d_ptm/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "# os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"of_3d_ptm/train/ar_model_latest.pth\")\n",
    "\n",
    "# print('Rename a trained model: ')\n",
    "# print('---------------------')\n",
    "# !cp $CHECKPOINT $HOST_RESULTS_DIR/of_3d_ptm/train/of_only_model.tlt\n",
    "# !ls -ltrh $HOST_RESULTS_DIR/of_3d_ptm/train/of_only_model.tlt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-4\"></a>\n",
    "\n",
    "We provide two different sample strategy to evaluate the pretrained model on video clips.\n",
    "\n",
    "* `center` mode: pick up the middle frames of a sequence to do inference. For example, if the model requires 32 frames as input and a video clip has 128 frames, then we will choose the frames from index 48 to index 79 to do the inference. \n",
    "* `conv` mode: convolutionly sample 10 sequences out of a single video and do inference. The final results are averaged."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate RGB model trained with PTM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model action_recognition evaluate \\\n",
    "                    -e $SPECS_DIR/experiment_rgb_3d_finetune.yaml \\\n",
    "                    encryption_key=$KEY \\\n",
    "                    results_dir=$RESULTS_DIR/rgb_3d_ptm \\\n",
    "                    dataset.workers=0 \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/rgb_3d_ptm/train/rgb_only_model.tlt  \\\n",
    "                    evaluate.batch_size=1 \\\n",
    "                    evaluate.test_dataset_dir=$DATA_DIR/test \\\n",
    "                    evaluate.video_eval_mode=center"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Optional:` Evaluate OF model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !tao model action_recognition evaluate \\\n",
    "#                     -e $SPECS_DIR/experiment_of_3d_finetune.yaml \\\n",
    "#                     encryption_key=$KEY \\\n",
    "#                     results_dir=$RESULTS_DIR/of_3d_ptm \\\n",
    "#                     dataset.workers=0 \\\n",
    "#                     evaluate.checkpoint=$RESULTS_DIR/of_3d_ptm/train/of_only_model.tlt  \\\n",
    "#                     evaluate.batch_size=1 \\\n",
    "#                     evaluate.test_dataset_dir=$DATA_DIR/test \\\n",
    "#                     evaluate.video_eval_mode=center"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-5\"></a>\n",
    "In this section, we run the action recognition inference tool to generate inferences with the trained RGB models and print the results. \n",
    "\n",
    "There are also two modes for inference just like evaluation: `center` mode and `conv` mode. And the final output will show each input sequence label in the videos like:\n",
    "`[video_sample_path] [labels list for sequences in the video sample]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model action_recognition inference \\\n",
    "                    -e $SPECS_DIR/experiment_rgb_3d_finetune.yaml \\\n",
    "                    encryption_key=$KEY \\\n",
    "                    results_dir=$RESULTS_DIR/rgb_3d_ptm \\\n",
    "                    dataset.workers=0 \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/rgb_3d_ptm/train/rgb_only_model.tlt \\\n",
    "                    inference.inference_dataset_dir=$DATA_DIR/test/ride_bike \\\n",
    "                    inference.video_inf_mode=center"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Optional:` Inference with OF-only model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !tao model action_recognition inference \\\n",
    "#                     -e $SPECS_DIR/experiment_of_3d_finetune.yaml \\\n",
    "#                     encryption_key=$KEY \\\n",
    "#                     results_dir=$RESULTS_DIR/of_3d_ptm \\\n",
    "#                     dataset.workers=0 \\\n",
    "#                     inference.checkpoint=$RESULTS_DIR/of_3d_ptm/train/of_only_model.tlt \\\n",
    "#                     inference.inference_dataset_dir=$DATA_DIR/test/ride_bike \\\n",
    "#                     inference.video_inf_mode=center"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy! <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Export the RGB model to encrypted ONNX model\n",
    "!tao model action_recognition export \\\n",
    "                   -e $SPECS_DIR/experiment_rgb_3d_finetune.yaml \\\n",
    "                   encryption_key=$KEY \\\n",
    "                   results_dir=$RESULTS_DIR/rgb_3d_ptm \\\n",
    "                   export.checkpoint=$RESULTS_DIR/rgb_3d_ptm/train/rgb_only_model.tlt \\\n",
    "                   export.onnx_file=$RESULTS_DIR/export/rgb_resnet18_3.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end. You may continue by deploying this RGB model to [DeepStream](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_3D_Action.html)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Optional` Export OF-Only model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Export the OF model to encrypted ONNX model\n",
    "# !tao model action_recognition export \\\n",
    "#                    -e $SPECS_DIR/experiment_of_3d_finetune.yaml \\\n",
    "#                    encryption_key=$KEY \\\n",
    "#                    results_dir=$RESULTS_DIR/of_3d_ptm \\\n",
    "#                    export.checkpoint=$RESULTS_DIR/of_3d_ptm/of_only_model.tlt\\\n",
    "#                    export.onnx_file=$RESULTS_DIR/export/of_resnet18_3.onnx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OF model is not supported in DeepStream. But you can play with stand-alone TensorRT inference in [tao_toolkit_recipes](https://github.com/NVIDIA-AI-IOT/tao_toolkit_recipes/tree/main/tao_action_recognition/tensorrt_inference) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
